/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/* !
 * \file auto_tiling.h
 * \brief
 */

#ifndef OPS_BUILT_IN_OP_TILING_AUTO_TILING_CONTEXT_H
#define OPS_BUILT_IN_OP_TILING_AUTO_TILING_CONTEXT_H

#include <iterator>

#include "graph/utils/op_desc_utils.h"
#include "exe_graph/runtime/tiling_context.h"
#include "exe_graph/lowering/shape_utils.h"
#include "exe_graph/lowering/getcdim.h"
#include "external/graph/operator.h"
#include "register/op_tiling_info.h"
#include "vector_op_info.h"
#include "op_tiling.h"
#include "vector_tiling_rt2.h"

namespace TbeReduce {
class OpShape {
public:
    OpShape() = default;
    explicit OpShape(const gert::Shape *_rt_shape) : rt_shape(_rt_shape){};
    explicit OpShape(const ge::GeShape *_ge_shape) : ge_shape(_ge_shape){};
    ~OpShape() = default;
    inline size_t GetDimNum() const
    {
        if (rt_shape != nullptr) {
            return rt_shape->GetDimNum();
        }
        return ge_shape->GetDimNum();
    };
    inline int64_t GetDim(size_t idx) const
    {
        if (rt_shape != nullptr) {
            return rt_shape->GetDim(idx);
        }
        return ge_shape->GetDim(idx);
    };
    inline int64_t GetShapeSize() const
    {
        if (rt_shape != nullptr) {
            return rt_shape->GetShapeSize();
        }
        return ge_shape->GetShapeSize();
    };
    inline bool Empty() const
    {
        return rt_shape != nullptr && ge_shape != nullptr;
    };

private:
    const gert::Shape *rt_shape{ nullptr };
    const ge::GeShape *ge_shape{ nullptr };
};

template <typename T1, typename T2>
static void GetDataToVector(const uint8_t *const_data, size_t data_size, std::vector<T1> &result)
{
    size_t size = data_size / sizeof(T2);
    result.resize(size);
    const T2 *data = reinterpret_cast<const T2 *>(const_data);
    for (size_t i = 0; i < size; i++) {
        result[i] = *(data + i);
    }
}

template <typename T>
bool GetConstIntData(const ge::Operator &paras, const int64_t const_input_idx, std::vector<T> &values)
{
    auto op_desc = ge::OpDescUtils::GetOpDescFromOperator(paras);
    ge::ConstGeTensorBarePtr const_tensor = ge::OpDescUtils::GetInputConstData(paras, const_input_idx);
    if (const_tensor == nullptr) {
        auto input_name = op_desc->GetInputNameByIndex(const_input_idx);
        HCCL_ERROR("GetConstIntData constvalue [%s] is not exists.", input_name.c_str());
        return false;
    }

    const auto &tensor_data = const_tensor->GetData();
    auto data = tensor_data.GetData();
    if (data == nullptr) {
        auto input_name = op_desc->GetInputNameByIndex(const_input_idx);
        HCCL_ERROR("GetConstIntData constvalue [%s] is nullptr.", input_name.c_str());
        return false;
    }
    auto size = tensor_data.GetSize();
    ge::DataType dtype = op_desc->MutableInputDesc(const_input_idx)->GetDataType();
    switch (dtype) {
        case ge::DT_UINT64:
            GetDataToVector<T, uint64_t>(data, size, values);
            break;
        case ge::DT_INT64:
            GetDataToVector<T, int64_t>(data, size, values);
            break;
        case ge::DT_UINT32:
            GetDataToVector<T, uint32_t>(data, size, values);
            break;
        case ge::DT_INT32:
            GetDataToVector<T, int32_t>(data, size, values);
            break;
        case ge::DT_FLOAT:
            GetDataToVector<T, float>(data, size, values);
            break;
        default: {
            HCCL_ERROR("GetConstIntData GetConstValue of dtype[%d] has not implement.", dtype);
            return false;
        }
    }
    return true;
}

template <typename T> bool GetData(const gert::Tensor *tensor, std::vector<int64_t> &values)
{
    if (tensor == nullptr) {
        return false;
    }
    size_t shape_size = tensor->GetShapeSize();
    values.resize(shape_size);
    auto *tensor_data = tensor->GetData<T>();
    if (tensor_data == nullptr) {
        return false;
    }
    for (size_t i = 0; i < shape_size; i++) {
        values[i] = static_cast<int64_t>(*(tensor_data + i));
    }
    return true;
}

static const OpShape empty_shape;

class AutoTilingOp {
public:
    explicit AutoTilingOp(const char *_op_type, const ge::Operator *_op_paras,
        const AutoTilingCompileInfo *_compile_info, optiling::utils::OpRunInfo *_run_info)
        : op_type(_op_type), op_paras(_op_paras), compile_info(_compile_info), run_info(_run_info)
    {
        op_desc_ptr = ge::OpDescUtils::GetOpDescFromOperator(*op_paras);
    };
    ~AutoTilingOp() = default;

public:
    bool OutputGetSubFormat(size_t idx, int32_t &groups);
    bool GetInputFormat(size_t idx, ge::Format &format);
    bool GetInputOriFormat(size_t idx, ge::Format &format);
    bool GetOutputFormat(size_t idx, ge::Format &format);
    bool GetInputDataType(size_t idx, ge::DataType &dtype)
    {
        if (op_desc_ptr == nullptr) {
            return false;
        }
        auto input_desc = op_desc_ptr->MutableInputDesc(idx);
        if (input_desc == nullptr) {
            return false;
        }
        dtype = input_desc->GetDataType();
        return true;
    }
    bool GetInputDataType(const OpInfoImpl *op_info, ge::DataType &dtype)
    {
        if (op_info && op_info->GetInType()) {
            dtype = *op_info->GetInType();
            return true;
        }
        return GetInputDataType(static_cast<size_t>(0), dtype);
    }
    bool GetOutputDataType(size_t idx, ge::DataType &dtype)
    {
        if (op_desc_ptr == nullptr) {
            return false;
        }
        auto output_desc = op_desc_ptr->MutableOutputDesc(idx);
        if (output_desc == nullptr) {
            return false;
        }
        dtype = output_desc->GetDataType();
        return true;
    }

    size_t GetInputNums()
    {
        if (op_desc_ptr == nullptr) {
            return 0;
        }
        return op_desc_ptr->GetInputsSize();
    }
    size_t GetInputNums(const OpInfoImpl *op_info)
    {
        if (op_info != nullptr && op_info->GetInputNum() > 0) {
            return op_info->GetInputNum();
        }
        return GetInputNums();
    }

    size_t GetOutputNums()
    {
        if (op_desc_ptr == nullptr) {
            return 0;
        }
        return op_desc_ptr->GetOutputsSize();
    }
    OpShape GetInputShape(size_t idx)
    {
        if (op_desc_ptr == nullptr) {
            return empty_shape;
        }
        auto input_desc = op_desc_ptr->MutableInputDesc(idx);
        if (input_desc == nullptr) {
            return empty_shape;
        }
        const ge::GeShape &shape = input_desc->GetShape();
        OpShape opShape(&shape);
        return opShape;
    }
    OpShape GetOriginInputShape(size_t idx);
    OpShape GetOutputShape(size_t idx)
    {
        if (op_desc_ptr == nullptr) {
            return empty_shape;
        }
        auto output_desc = op_desc_ptr->MutableOutputDesc(idx);
        if (output_desc == nullptr) {
            return empty_shape;
        }
        const ge::GeShape &shape = output_desc->GetShape();
        OpShape opShape(&shape);
        return opShape;
    }
    OpShape GetOriginOutputShape(size_t idx);
    OpShape GetInputShape(const OpInfoImpl *op_info, size_t idx)
    {
        if (op_info != nullptr && idx < op_info->GetInputNum()) {
            auto shapes = op_info->GetInputShapeVector();
            OpShape opShape(&(shapes[idx]));
            return opShape;
        }
        return GetInputShape(idx);
    }
    OpShape GetOutputShape(const OpInfoImpl *op_info, size_t idx);
    const char *GetOpType()
    {
        return op_type;
    }
    const char *GetOpName()
    {
        return op_type;
    }
    const AutoTilingCompileInfo *GetCompileInfo()
    {
        return compile_info;
    }
    bool SetBlockDim(uint32_t block_dims)
    {
        run_info->SetBlockDim(block_dims);
        return true;
    }
    bool SetTilingKey(uint64_t tiling_key)
    {
        run_info->SetTilingKey(tiling_key);
        return true;
    }
    bool SetNeedAtomic(bool flag);
    void SetCompileInfo(const AutoTilingCompileInfo *_compile_info);
    bool GetAttr(const char *name, size_t index, bool &value);
    bool GetAttr(const char *name, size_t index, int64_t &value);
    bool GetAttr(const char *name, size_t index, std::vector<int64_t> &values);
    bool GetAttr(const char *name, size_t index, std::string &value);
    bool GetConstInput(const char *name, size_t index, int64_t &value)
    {
        if (name != nullptr) {
            index = op_desc_ptr->GetInputIndexByName("axes");
        }
        std::vector<int64_t> values{ 0 };
        if (!GetConstIntData(*op_paras, index, values)) {
            return false;
        }
        if (values.empty()) {
            return false;
        }
        value = values[0];
        return true;
    }
    bool GetConstInput(const char *name, size_t index, std::vector<int64_t> &values)
    {
        constexpr size_t max_integer = 2147483647;
        if (name != nullptr && index == max_integer) {
            index = op_desc_ptr->GetInputIndexByName("axes");
        }
        if (!GetConstIntData(*op_paras, index, values)) {
            return false;
        }
        return true;
    }
    bool WriteVarAttrs(const uint64_t tiling_key) const
    {
        if (compile_info->var_attr_wrap.IsEmpty()) {
            return true;
        }
        return compile_info->var_attr_wrap.WriteVarAttrs(tiling_key, op_type, *op_paras, *run_info);
    }
    const ge::Operator *GetOpParas();
    optiling::utils::OpRunInfo *GetRunInfo();
    int64_t GetOriginalInputDimC(const size_t index)
    {
        // FE not support runtime1.0 fetch dim c
        return -1;
    }

    template <typename ForwardIterator> bool AddWorkspace(ForwardIterator first, size_t n)
    {
        for (size_t i = 0; i < n; i++) {
            run_info->AddWorkspace(static_cast<int64_t>(*first));
            first++;
        }
        return true;
    }

    template <typename T> bool Append(const T &data)
    {
        run_info->AddTilingData(data);
        return true;
    }

    template <typename T> bool Append(const T *data, size_t tiling_num)
    {
        run_info->AddTilingData((char *)data, tiling_num * sizeof(T));
        return true;
    }

private:
    const char *op_type{ nullptr };
    const ge::Operator *op_paras{ nullptr };
    const AutoTilingCompileInfo *compile_info{ nullptr };
    optiling::utils::OpRunInfo *run_info{ nullptr };
    ge::OpDescPtr op_desc_ptr{ nullptr };
};

class AutoTilingContext {
public:
    explicit AutoTilingContext(gert::TilingContext *_context) : context(_context)
    {
        tiling_data = context->GetRawTilingData();
    }
    explicit AutoTilingContext(gert::TilingContext *_context, const AutoTilingCompileInfo *_compile_info)
        : context(_context), compile_info(_compile_info)
    {
        tiling_data = context->GetRawTilingData();
    }
    ~AutoTilingContext() = default;

public:
    bool OutputGetSubFormat(size_t idx, int32_t &groups);
    bool GetInputFormat(size_t idx, ge::Format &format);
    bool GetInputOriFormat(size_t idx, ge::Format &format);
    bool GetOutputFormat(size_t idx, ge::Format &format);
    bool GetInputDataType(size_t idx, ge::DataType &dtype)
    {
        auto input_desc = context->GetInputDesc(idx);
        if (input_desc == nullptr) {
            return false;
        }
        dtype = input_desc->GetDataType();
        return true;
    }
    bool GetInputDataType(const OpInfoImpl *op_info, ge::DataType &dtype)
    {
        if (op_info && op_info->GetInType()) {
            dtype = *op_info->GetInType();
            return true;
        }
        return GetInputDataType(static_cast<size_t>(0), dtype);
    }
    bool GetOutputDataType(size_t idx, ge::DataType &dtype)
    {
        auto output_desc = context->GetOutputDesc(idx);
        if (output_desc == nullptr) {
            return false;
        }
        dtype = output_desc->GetDataType();
        return true;
    }
    size_t GetInputNums()
    {
        return context->GetComputeNodeInputNum();
    }
    size_t GetInputNums(const OpInfoImpl *op_info)
    {
        if (op_info != nullptr && op_info->GetInputNum() > 0) {
            HCCL_DEBUG("Get custom input shape num success");
            return op_info->GetInputNum();
        }
        return GetInputNums();
    }

    size_t GetOutputNums()
    {
        auto compute_node_info = context->GetComputeNodeInfo();
        if (compute_node_info == nullptr) {
            return 0;
        }
        return compute_node_info->GetOutputsNum();
    }
    OpShape GetInputShape(size_t idx)
    {
        auto input_shape = context->GetInputShape(idx);
        if (input_shape == nullptr) {
            return empty_shape;
        }
        const gert::Shape &shape = gert::EnsureNotScalar(input_shape->GetStorageShape());
        OpShape opShape(&shape);
        return opShape;
    }
    OpShape GetOriginInputShape(size_t idx);
    OpShape GetOutputShape(size_t idx)
    {
        const gert::Shape &shape = gert::EnsureNotScalar(context->GetOutputShape(idx)->GetStorageShape());
        OpShape opShape(&shape);
        return opShape;
    }
    OpShape GetOriginOutputShape(size_t idx);
    OpShape GetInputShape(const OpInfoImpl *op_info, size_t idx)
    {
        if (op_info != nullptr && idx < op_info->GetInputNum()) {
            auto shapes = op_info->GetInputShapeVector();
            OpShape opShape(&(shapes[idx]));
            return opShape;
        }
        return GetInputShape(idx);
    }
    OpShape GetOutputShape(const OpInfoImpl *op_info, size_t idx);
    const char *GetOpType()
    {
        return context->GetNodeType();
    }
    const char *GetOpName()
    {
        return context->GetNodeName();
    }
    const AutoTilingCompileInfo *GetCompileInfo()
    {
        if (compile_info != nullptr) {
            return compile_info;
        }
        return reinterpret_cast<const AutoTilingCompileInfo *>(context->GetCompileInfo());
    }
    bool SetBlockDim(uint32_t block_dims)
    {
        if (context->SetBlockDim(block_dims) == ge::GRAPH_FAILED) {
            return false;
        }
        return true;
    }
    bool SetTilingKey(uint64_t tiling_key)
    {
        if (context->SetTilingKey(tiling_key) == ge::GRAPH_FAILED) {
            return false;
        }
        return true;
    }
    bool SetNeedAtomic(bool flag);
    void SetCompileInfo(const AutoTilingCompileInfo *_compile_info)
    {
        compile_info = _compile_info;
    }
    bool GetAttr(const char *name, size_t index, bool &value);
    bool GetAttr(const char *name, size_t index, int64_t &value);
    bool GetAttr(const char *name, size_t index, std::vector<int64_t> &values);
    bool GetAttr(const char *name, size_t index, std::string &value);
    bool GetAttr(const char *name, size_t index, float &value);
    bool GetConstInput(const char *name, size_t index, int64_t &value)
    {
        std::vector<int64_t> const_values;
        if (GetConstInput(name, index, const_values)) {
            if (!const_values.empty()) {
                value = const_values[0];
                return true;
            }
        }
        return false;
    }
    bool GetConstInput(const char *name, size_t index, std::vector<int64_t> &values)
    {
        auto *tensor = context->GetInputTensor(index);
        if (tensor == nullptr) {
            return false;
        }
        if (tensor->GetDataType() == ge::DT_INT32) {
            return GetData<int32_t>(tensor, values);
        } else if (tensor->GetDataType() == ge::DT_INT64) {
            return GetData<int64_t>(tensor, values);
        } else if (tensor->GetDataType() == ge::DT_UINT32) {
            return GetData<uint32_t>(tensor, values);
        } else if (tensor->GetDataType() == ge::DT_UINT64) {
            return GetData<uint64_t>(tensor, values);
        }
        HCCL_ERROR("Get const input error, unsupported data type.");
        return false;
    }
    bool WriteVarAttrs(const uint64_t tiling_key)
    {
        const AutoTilingCompileInfo *autoTilingCompileInfo =
            static_cast<const AutoTilingCompileInfo *>(GetCompileInfo());
        if (autoTilingCompileInfo->var_attr_wrap.IsEmpty()) {
            return true;
        }
        return autoTilingCompileInfo->var_attr_wrap.WriteVarAttrs(tiling_key, *context);
    }
    const ge::Operator *GetOpParas();
    optiling::utils::OpRunInfo *GetRunInfo();
    int64_t GetOriginalInputDimC(const size_t index)
    {
        return gert::GetInputCDim(context, index);
    }

    template <typename ForwardIterator> bool AddWorkspace(ForwardIterator first, size_t n)
    {
        auto workspace_data = context->GetWorkspaceSizes(n);
        if (workspace_data == nullptr) {
            return false;
        }
        for (size_t i = 0; i < n; i++) {
            workspace_data[i] = *first;
            first++;
        }
        return true;
    }

    template <typename T> bool Append(const T &data)
    {
        if (tiling_data->Append(data) == ge::GRAPH_FAILED) {
            return false;
        }
        return true;
    }

    template <typename T> bool Append(const T *data, size_t tiling_num)
    {
        if (tiling_data->Append(data, tiling_num) == ge::GRAPH_FAILED) {
            return false;
        }
        return true;
    }

private:
    gert::TilingContext *context{ nullptr };
    gert::TilingData *tiling_data{ nullptr };
    const AutoTilingCompileInfo *compile_info{ nullptr };
};
} // namespace optiling

#endif // OPS_BUILT_IN_OP_TILING_AUTO_TILING_CONTEXT_H
